#!/usr/bin/python
# -*- encoding:utf-8 -*-
import os
import pywt
import numpy as np
import matplotlib
matplotlib.use('agg')
import pylab as plt

# from utils.feature import load_audio


wavelet = 'db2'
level = 4
order = "freq"  # other option is "normal"
interpolation = 'nearest'
cmap = plt.cm.cool


def extract_fbank(wav):
    # data = load_audio(wav, mono=True, fs=44100)
    x = np.linspace(0, 1, num=512)
    data = np.sin(250 * np.pi * x**2)
    
    # Construct wavelet packet
    wp = pywt.WaveletPacket(data, wavelet, 'symmetric', maxlevel=level)
    nodes = wp.get_level(level, order=order)
    labels = [n.path for n in nodes]
    values = np.array([n.data for n in nodes], 'd')
    values = abs(values)

    # Show signal and wavelet packet coefficients
    fig = plt.figure()
    fig.subplots_adjust(hspace=0.2, bottom=.03, left=.07, right=.97, top=.92)
    ax = fig.add_subplot(2, 1, 1)
    ax.set_title("linchirp signal")
    ax.plot(x, data, 'b')
    ax.set_xlim(0, x[-1])

    ax = fig.add_subplot(2, 1, 2)
    ax.set_title("Wavelet packet coefficients at level %d" % level)
    ax.imshow(values, interpolation=interpolation, cmap=cmap, aspect="auto",
              origin="lower", extent=[0, 1, 0, len(values)])
    ax.set_yticks(np.arange(0.5, len(labels) + 0.5), labels)

    # Show spectrogram and wavelet packet coefficients
    fig2 = plt.figure()
    ax2 = fig2.add_subplot(211)
    ax2.specgram(data, NFFT=64, noverlap=32, Fs=2, cmap=cmap,
                interpolation='bilinear')
    ax2.set_title("Spectrogram of signal")
    ax3 = fig2.add_subplot(212)
    ax3.imshow(values, origin='upper', extent=[-1, 1, -1, 1],
            interpolation='nearest')
    ax3.set_title("Wavelet packet coefficients")
    plt.show()
    fig.savefig('test.png')

def main():
    extract_fbank('/home/xiaorong/Data/SED_SMALL/fire/kiva.wav')


if __name__ == '__main__':
    main()